from pydantic import BaseModel
from typing import List, Optional

class ChatContext(BaseModel):
    history: List[str]

class ClarifyPromptInput(BaseModel):
    botId: str
    siteId: int
    llm: str = 'GPT3.5'
    query: str
    context: List[ChatContext]

class QueryPromptInput(BaseModel):
    botId: str
    siteId: int
    llm: Optional[str] = 'GPT3.5'
    query: str
    standAloneQuestions: List[str]
    context: List[ChatContext]

class Prompt(BaseModel):
    role: str
    content: str

    def __init__(self, role: str, content: str) -> None:
        super().__init__(role=role, content=content)

class Article(BaseModel):
    number: int
    title: str
    url: str
    score: float

    def __init__(self, number: int, title: str, url: str, score: float) -> None:
        super().__init__(number=number, title=title, url=url, score=score)

class ClarifyPromptOutput(BaseModel):
    code: int = 0
    msg: str | None
    type: str = 'noAnswer'
    prompt: List[Prompt] = []

    def __init__(self, code: int =0, msg: str = None, type: str = 'noAnswer', prompt: List[Prompt] = []) -> None:
        super().__init__(code=code, msg=msg, type=type, prompt=prompt)

class QueryPromptOutput(BaseModel):
    code: int =0
    msg: str | None
    prompt: List[Prompt] = None
    articles: List[Article] = []

    def __init__(self, code: int =0, msg: str = None, prompt: List[Prompt]= [], articles: List[Article]= []) -> None:
        super().__init__(code=code, msg=msg, prompt=prompt, articles=articles)

class QueryInput(BaseModel):
    prompt: List[Prompt] = None
    articles: List[Article] = []

    def __init__(self, prompt: List[Prompt]= [], articles: List[Article]= []) -> None:
        super().__init__(prompt=prompt, articles=articles)

class QueryOutput(BaseModel):
    code: int =0
    msg: str | None
    prompt: List[Prompt] = None
    articles: List[Article] = []

    def __init__(self, code: int =0, msg: str = None, prompt: List[Prompt]= [], articles: List[Article]= []) -> None:
        super().__init__(code=code, msg=msg, prompt=prompt, articles=articles)